import requests
import json
from typing import Optional, Dict, Any
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelAPIClient:
    def __init__(self, base_url: str = "http://localhost:8000"):
        """
        初始化API客户端
        
        :param base_url: API基础地址，默认本地8000端口
        """
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        self.session.headers.update({
            "Content-Type": "application/json",
            "Accept": "application/json"
        })
        logger.info(f"初始化API客户端，服务地址: {self.base_url}")

    def generate_text(
        self,
        prompt: str,
        max_length: Optional[int] = 512,
        temperature: Optional[float] = 0.7,
        top_p: Optional[float] = 0.9,
        top_k: Optional[int] = 50,
        repetition_penalty: Optional[float] = 1.1,
        timeout: int = 60
    ) -> str:
        """
        生成文本
        
        :param prompt: 输入的提示文本
        :param max_length: 生成的最大长度
        :param temperature: 温度参数
        :param top_p: 核心采样参数
        :param top_k: 顶部k采样参数
        :param repetition_penalty: 重复惩罚系数
        :param timeout: 请求超时时间(秒)
        :return: 生成的文本
        :raises: Exception 当请求失败时抛出异常
        """
        endpoint = f"{self.base_url}/generate"
        payload = {
            "prompt": prompt,
            "max_length": max_length,
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "repetition_penalty": repetition_penalty
        }
        
        try:
            logger.info(f"发送生成请求，提示长度: {len(prompt)}字符")
            response = self.session.post(
                endpoint,
                data=json.dumps(payload),
                timeout=timeout
            )
            
            # 检查响应状态
            response.raise_for_status()
            
            # 解析响应
            result = response.json()
            generated_text = result.get("generated_text", "")
            
            logger.info(f"生成完成，生成长度: {len(generated_text)}字符")
            return generated_text
            
        except requests.exceptions.RequestException as e:
            error_msg = f"API请求失败: {str(e)}"
            if hasattr(e, 'response') and e.response is not None:
                try:
                    error_detail = e.response.json().get("detail", "无详细信息")
                    error_msg += f" | 服务端错误: {error_detail}"
                except ValueError:
                    error_msg += f" | 状态码: {e.response.status_code}"
            logger.error(error_msg)
            raise Exception(error_msg)
        except json.JSONDecodeError as e:
            error_msg = f"响应解析失败: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)
        except Exception as e:
            error_msg = f"未知错误: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)

    def health_check(self) -> Dict[str, Any]:
        """
        检查服务健康状态
        
        :return: 健康状态字典
        :raises: Exception 当请求失败时抛出异常
        """
        endpoint = f"{self.base_url}/health"
        
        try:
            response = self.session.get(endpoint, timeout=5)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            error_msg = f"健康检查失败: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)

# 使用示例
if __name__ == "__main__":
    # 1. 初始化客户端
    client = ModelAPIClient(base_url = "http://localhost:8000")  # 默认本地地址
    
    # 2. 检查服务健康状态
    try:
        health = client.health_check()
        print(f"服务状态: {health}")
    except Exception as e:
        print(f"服务不可用: {e}")
        exit(1)
    
    # 3. 生成文本示例
    try:
        prompt = "你好，请介绍一下你自己"
        generated_text = client.generate_text(
            prompt=prompt,
            max_length=200,
            temperature=0.7
        )
        print("\n生成结果:")
        print("-" * 50)
        print(generated_text)
        print("-" * 50)
    except Exception as e:
        print(f"生成失败: {e}")